# -*- coding: utf-8 -*-
"""
Created on 2018 3.30
@author: hugh
"""

import argparse
import time
import os
import config
from load_image import LoadImage
from sklearn.externals import joblib
import cv2

parser = argparse.ArgumentParser()
parser.add_argument("-p", "--image_path", default="./test_dir/bush_powell.jpg",
						help="path of the image needed to be recognition")
args = parser.parse_args()

def recognition(X_test, model_check_point=config.model_check_point):
	"""识别函数
	Args：
		X_test: 测试图片
		model_check_point：模型存放路径
	Return：识别出的标签list
	"""
	# 加载KNN Classifier模型	
	model = joblib.load(model_check_point)
	#进行识别
	predict = model.predict(X_test)
	return predict.tolist()


if __name__ == '__main__':
	start = time.time()
	#检测输入的图片是否存在
	if not os.path.exists(args.image_path):
		print("Image is not exists.")
		exit(1)
	#检测训练的模型和文件是否存在
	if not os.path.exists(config.train_label_name_data) or \
		not os.path.exists(config.train_face_descriptor_data) or \
		not os.path.exists(config.model_check_point):
		print("model file is not exists, please train firstly")
		print("python train.py")
		exit(1)
	print ("-----------------------------load_image.py start---------------------")
	# 加载数据
	all_image = LoadImage(config.data_dir)
	train_label_name_dict = all_image.train_label_name_dict
	#检测人脸，识别人脸并转换成128维特征向量，组合成预测数据
	X_test, face_dets = all_image.get_feture_vector_from_image(args.image_path)
	#模型预测人脸，获取人脸对应的标签
	predicts = recognition(X_test)
	#将人脸标签转换为人名输出
	for i, predict in enumerate(predicts):
		print("Face:{},Recognition:{}".format(i, train_label_name_dict.get(predict)))
	# 读取测试图片
	img = cv2.imread(args.image_path)
	for i, d in enumerate(face_dets):
		#给人脸标记矩形框
		img = cv2.rectangle(img,(d.left(), d.top()), (d.right(), d.bottom()),(0, 255, 0), 1)
		#给人脸标记姓名
		cv2.putText(img, train_label_name_dict.get(predicts[i]), (d.left(), d.top() - 5),\
											cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
	# 显示标记后的测试图片
	cv2.imshow('Face recognition result', img)
	# 计算整个识别所耗时间
	end = time.time()
	print("Recognition end. Elapsed time:{}".format(end - start))
	# 等待按任意键
	cv2.waitKey()
	cv2.destroyAllWindows()
